#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2021/11/1 14:37
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : naive_bayes_classifier.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
from sklearn.naive_bayes import GaussianNB


def load_data(data_path, delimiter=None, attr_col=None, label_col=-1, reindex=None, test_size=0.2):
    """
    加载数据集并且划分特征集和标记集
    :param data_path: 数据集文件路径
    :param delimiter: 分隔符
    :param attr_col: 属性列索引
    :param label_col: 标记列索引
    :param reindex: 重定义列索引，元组
    :param test_size: 划分测试集比例
    :return:
    """
    import pandas as pd
    from sklearn.model_selection import train_test_split
    dataset = pd.read_csv(data_path, delimiter=delimiter, names=list(range(*reindex)))
    attr_data, label_data = dataset.iloc[:, :attr_col], dataset.iloc[:, label_col]
    return train_test_split(attr_data, label_data, test_size=test_size, random_state=1)


def train(feature_train, target_train, naive_bayes_type=GaussianNB):
    """
    训练模型
    :param feature_train: 特征训练集
    :param target_train: 目标训练集
    :param naive_bayes_type: 贝叶斯分类器类型
    :return: 训练完成的模型
    """
    classifier_model = naive_bayes_type()  # assume the distribution is gaussian
    classifier_model.fit(feature_train, target_train)
    return classifier_model


def evaluate(naive_bayes_classifier, feature_test, target_test):
    """
    评估模型的精度，查准率，召回率
    :param naive_bayes_classifier: 朴素贝叶斯分类器
    :param feature_test: 测试特征集
    :param target_test: 测试标记集
    :return: 精度，查准率，召回率
    """
    from sklearn.metrics import accuracy_score, precision_score, recall_score
    target_pred = naive_bayes_classifier.predict(feature_test)
    target = target_test, target_pred
    return accuracy_score(*target), precision_score(*target, average=None), recall_score(*target, average=None)


if __name__ == '__main__':
    # 加载并划分数据集
    DATA_PATH = "../dataset/shuiguodata.csv"
    attr_train, attr_test, label_train, label_test = load_data(DATA_PATH)
    naive_bayes_dict = {"高斯朴素贝叶斯": GaussianNB}
    for nb_type, value in naive_bayes_dict.items():
        model = train(attr_train, label_train)
        print("{}分类器\n精度:{}\n每一类查准率:{}\n每一类召回率:{}".format(
            nb_type, *evaluate(model, attr_test, label_test)))
